"""
-*- coding: utf-8 -*-
@File  : generate_annotation.py
@author: ZhenyuYang
@Time  : 2023/04/05 21:37
"""
import pandas as pd
import math
import os

dataset = ["casme2", "smic", "samm"]
cross = 5

annotation_path = "../combined_3class_gt.csv"
annotations_frame = pd.read_csv(annotation_path, header=None)
annotations = annotations_frame.to_numpy()
if len(dataset) == 3:
    save_path = "../data/annotations"
elif len(dataset) == 1:
    save_path = "../data/{}_annotations".format(dataset[0])
else:
    print("Dataset out of support.")
    raise Exception
if not os.path.exists(save_path):
    os.mkdir(save_path)
sub_list = []
for sub in annotations:
    sub_name = "{}-{}".format(sub[0], sub[1])
    if sub[0] in dataset and sub_name not in sub_list:
        sub_list.append(sub_name)
print("Subject num: {}".format(len(sub_list)))

labels = ["negative", "positive", "surprise"]
with open("{}/classInd.txt".format(save_path), "w") as f:
    for label_index, label in enumerate(labels):
        f.write("{} {}\n".format(label_index+1, label))

test_num = math.floor(len(sub_list) / cross)
for cross_index in range(cross):
    test_start_index = cross_index * test_num
    test_end_index = min((cross_index + 1) * test_num, len(sub_list))
    # print(test_start_index, test_end_index)
    train_sub = []
    test_sub = []
    for sub_index, sub in enumerate(sub_list):
        if (sub_index >= test_start_index) & (sub_index < test_end_index):
            test_sub.append(sub)
        else:
            train_sub.append(sub)
    print("Cross {}: \ntrain: {}\ntest: {}".format(cross_index+1, train_sub, test_sub))
    with open("{}/trainlist{:02d}.txt".format(save_path, cross_index+1), "w") as f:
        pass
    with open("{}/testlist{:02d}.txt".format(save_path, cross_index+1), "w") as f:
        pass
    for label in labels:
        data_list = os.listdir("../data/videos/{}".format(label))
        for data in data_list:
            sub_info = data.split("-")
            if sub_info[0] not in dataset:
                continue
            sub_name = "{}-{}".format(sub_info[0], sub_info[1])
            if sub_name in train_sub:
                with open("{}/trainlist{:02d}.txt".format(save_path, cross_index + 1), "a") as f:
                    f.write("{}/{}\n".format(label, data))
            elif sub_name in test_sub:
                with open("{}/testlist{:02d}.txt".format(save_path, cross_index + 1), "a") as f:
                    f.write("{}/{}\n".format(label, data))
            else:
                print("Error: cross {}, sub {}".format(cross_index, sub_name))

